#!/usr/bin/env python
from racecar.bayes import *
from racecar.parse import *
from racecar.data import *
from racecar.cluster import *
import os

class DataTests:
  def test1():
    # Extract data.
    data = parseBayesData(r"..\..\data\tests\test6.tab");
    
    # Print data back out.
    print "--> Data load and write test..";
    data.write();
    print "";
  
  test1 = staticmethod(test1);

class PosteriorExtractorTests:
  # Test ze parsing.
  def test1():
    ext = PosteriorExtractor();
    ext.parse(r'..\data\output\test.posterior.tab');
    ext.parse(r'..\data\output\test2.posterior.tab');
    
    print "--> Posterior extractor test..";
    ext.getPosterior().write();
    print "";
    
  # Test parsing directly from program output.
  def test2():
    ext = PosteriorExtractor();
    f = os.popen(r"..\..\bin\Racecar.exe -i ..\..\data\tests\test4.txt --ladder --cutoffs 50,50 --start_k 1 -s 10 -n 0");
    ext.parseFileObject(f);
    
    print "--> Posterior extraction from stdout of racecar..";
    ext.getPosterior().write();
    print "";
  
  test1 = staticmethod(test1);
  test2 = staticmethod(test2);
  
class HungryPosteriorTests:
  # Test base functionality. 
  def test1():
    print "--> HungryPosterior base test..";
    
    # Our posterior.
    posterior = HungryPosterior();
    
    # Two networks, identical edges, different scores.
    a = Network();
    a.addEdge(1,0);
    a.addEdge(2,4);
    a.addEdge(3,5);
    b = Network();
    b.addEdge(1,0);
    b.addEdge(2,4);
    b.addEdge(3,5);
    
    posterior.update(-1,a);
    posterior.update(-1,b);
    
    # Some random network/score.
    c = Network();
    c.addEdge(2,7);
    posterior.update(-4,c);
    
    # Print 'em!
    posterior.write();
    print "";

  # Test score iteration.
  def test2():
    post = HungryPosterior();
    
    a = Network();
    b = Network();
    c = Network();
    d = Network();
    
    a.addEdge(1,0);
    b.addEdge(3,0);
    c.addEdge(4,1);
    d.addEdge(5,7);
    
    post.update(-1,a);
    post.update(-2.7,b);
    post.update(-4,c);
    post.update(-1.43,d);
    
    print "--> HungryPosterior score iterator test..";
    print [i for i in post.scores()]; 
    print "";
  
  test1 = staticmethod(test1);
  test2 = staticmethod(test2);

class UniquePosteriorTests:
  # Test getting scores.
  def test1():
    post = UniquePosterior();
    
    a = Network();
    b = Network();
    c = Network();
    d = Network();
    
    a.addEdge(1,0);
    b.addEdge(3,0);
    c.addEdge(4,1);
    d.addEdge(5,7);
    
    post.update(-1,a);
    post.update(-2.7,b);
    post.update(-4,c);
    post.update(-1.43,d);
    
    print "--> UniquePosterior score iterator test..";
    print [i for i in post.scores()]; 
    print "";
    
  # Test iterating over networks.
  def test2():
    post = UniquePosterior();
    
    a = Network();
    b = Network();
    c = Network();
    d = Network();
    
    a.addEdge(1,0);
    b.addEdge(3,0);
    c.addEdge(4,1);
    d.addEdge(5,7);
    
    post.update(-1,a);
    post.update(-2.7,b);
    post.update(-4,c);
    post.update(-1.43,d);
    
    print "--> UniquePosterior network iterator test..";
    print [str(i) for i in post.networks()]; 
    print "";
  
  # Test pretty printing. 
  def test3():
    post = UniquePosterior();
    
    a = Network();
    b = Network();
    c = Network();
    d = Network();
    e = Network();
    
    a.addEdge(1,0);
    b.addEdge(3,0);
    b.addEdge(2,3);
    c.addEdge(4,1);
    d.addEdge(5,7);
    e.addEdge(4,3);
    
    post.update(-1,a);
    post.update(-2.7,b);
    post.update(-4,c);
    post.update(-1.43,d);
    post.update(-2.7,e);
    
    print "--> UniquePosterior print test..";
    post.write();
    print "";
    
  # Test iterating over network/score pairs. 
  def test4():
    post = UniquePosterior();
    
    a = Network();
    b = Network();
    c = Network();
    d = Network();
    
    a.addEdge(1,0);
    b.addEdge(3,0);
    c.addEdge(4,1);
    d.addEdge(5,7);
    
    post.update(-1,a);
    post.update(-2.7,b);
    post.update(-4,c);
    post.update(-1.43,d);
    
    print "--> UniquePosterior pair iterator test..";
    for score,net in post.pairs():
      print "Network: " + str(net) + "\tScore: " + str(score);
      
    print "";
    
  test1 = staticmethod(test1);
  test2 = staticmethod(test2);
  test3 = staticmethod(test3);
  test4 = staticmethod(test4);

class NetworkTests:
  # Equality test. 
  def test1():
    a = Network();
    b = Network();
    
    a.addEdge(1,0);
    a.addEdge(2,0);
    a.addEdge(4,-1);
    
    b.addEdge(1,0);
    b.addEdge(2,0);
    b.addEdge(4,-1);
    
    print "--> Network equality test..";
    print "Network A: " + str(a);
    print "Network B: " + str(b);
    print "A == B: " + str(a == b);
    print "";
  
  # Hash code test.
  def test2():
    a = Network();
    b = Network();
    c = Network();
    
    a.addEdge(1,0);
    a.addEdge(2,0);
    a.addEdge(4,-1);
    
    b.addEdge(1,0);
    b.addEdge(2,0);
    b.addEdge(4,-1);
    
    c.addEdge(4,-2);
    c.addEdge(5,1);
    
    print "--> Network hash code test..";
    print "Network A: " + str(a);
    print "Network B: " + str(b);
    print "Network C: " + str(c);
    print "Hash A: " + str(hash(a));
    print "Hash B: " + str(hash(b));
    print "Hash C: " + str(hash(c));
    print "";
  
  # Creating a network from string.
  def test3():
    a = Network.fromString("-1: 4, 0: 1 2");
    
    print "--> Network from string test..";
    print "String: -1: 4, 0: 1 2";
    print "Network: " + str(a);
    print "";
    
  # Static method declarations. 
  test1 = staticmethod(test1);
  test2 = staticmethod(test2);
  test3 = staticmethod(test3);
  
class ExactClusteratorTests:
  def test1():
    print "--> Clusterator test 1..";
        
    # Extract data.
    data = parseBayesData(r"..\..\data\tests\equivTest1.txt");

    # Print the data.
    data.write();
    print "";
    
    # Cluster it.
    cobj = ExactClusterator(data);
    cobj.cluster();
    cobj.printClusters();
    print "";
    
  def test2():
    print "--> Clusterator test 2..";
    
    # Extract data.
    data = parseBayesData(r"..\..\data\tests\test3.txt");
    
    # Print the data.
    data.write();
    print "";
    
    # Cluster it.
    cobj = ExactClusterator(data);
    cobj.cluster();
    cobj.printClusters();
    print "";
    
  def test3():
    print "--> Clusterator massive test..";
    
    # Extract data.
    print "Parsing data.."
    data = parseBayesData(r"..\..\data\input\genexp_imputed_log2_top20000var_response_ready.txt");
    
    # Print the data.
    #data.write();
    #print "";
    
    # Cluster it.
    print "Clustering.."
    cobj = ExactClusterator(data);
    cobj.cluster();
    cobj.printClusters();
    print "";
    
  def test4():
    print "--> Clusterator write test..";
        
    # Extract data.
    print "Parsing.."
    data = parseBayesData(r"..\..\data\tests\equivTest1.txt");

    # Print the data.
    data.write();
    print "";
    
    # Cluster it.
    cobj = ExactClusterator(data);
    print "Clustering.."
    cobj.cluster();
    print "Writing.."
    cobj.writeData(r"..\..\data\tests\equivTest1_clustered.tab",verbose=False);
    print "";
    
  test1 = staticmethod(test1);
  test2 = staticmethod(test2);
  test3 = staticmethod(test3);
  test4 = staticmethod(test4);

# Execute desired tests. 
#NetworkTests.test1();
#NetworkTests.test2();
#NetworkTests.test3();
#UniquePosteriorTests.test1();
#UniquePosteriorTests.test2();
#UniquePosteriorTests.test3();
#UniquePosteriorTests.test4();
#PosteriorExtractorTests.test1();
#PosteriorExtractorTests.test2();
#DataTests.test1();
#HungryPosteriorTests.test1();
#HungryPosteriorTests.test2();
ExactClusteratorTests.test1();
ExactClusteratorTests.test2();
#ExactClusteratorTests.test3();
ExactClusteratorTests.test4();